use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
use std::time::Duration;
use tokio::io;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, UdpSocket};

pub async fn run_proxy() -> io::Result<()> {
    let addr = "127.0.0.1:6789";
    tracing::info!("代理服务[{addr}]启动...");
    let listener = TcpListener::bind(addr).await?;

    loop {
        let (stream, _) = listener.accept().await?;
        // 处理客户端连接
        tokio::spawn(async move {
            if let Err(e) = handle_client2(stream).await {
                // eprintln!("Error handling client: {}", e);
            }
        });
    }
}
async fn handle_client2(mut stream: TcpStream) -> io::Result<()> {
    // println!("握手阶段@handle_client2: {:?}", stream);
    // 1. 握手阶段
    let mut buffer = [0u8; 256];
    stream.read(&mut buffer).await?;
    let version = buffer[0];
    let nmethods = buffer[1] as usize;
    // println!(
    //     "握手阶段 buffer: {:?}, version: {}, nmethods: {:?}",
    //     buffer, version, nmethods
    // );
    if version != 0x05 {
        tracing::error!("不支持的SOCKS版本: {:?}", version);
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "不支持的 SOCKS 版本",
        ));
    }
    let methods = &buffer[2..2 + nmethods];
    // println!("methods: {:?}", methods);
    let selected_method = if methods.contains(&0x00) {
        0x00 // 无认证
    } else {
        tracing::error!("不支持的认证方法: {:?}", methods);
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "不支持的认证方法",
        ));
    };
    // println!("selected_method: {:?}", selected_method);
    stream.write_all(&[0x05, selected_method]).await?;

    // 2.请求阶段
    // println!(
    //     "请求阶段@handle_client2: selected_method: {:?}",
    //     selected_method
    // );
    let n = stream.read(&mut buffer).await?;
    if n < 4 {
        tracing::error!("无效的请求数据: {:?}", buffer);
        return Err(io::Error::new(io::ErrorKind::InvalidData, "无效的请求数据"));
    }
    let version = buffer[0];
    let cmd = buffer[1];
    let _rsv = buffer[2]; // RSV
    let addr_type = buffer[3];
    // println!(
    //     "buffer: {:?}, version: {}, cmd: {}, _rsv: {}, addr_type: {}",
    //     buffer, version, cmd, _rsv, addr_type
    // );
    if version != 0x05 || (cmd != 0x01 && cmd != 0x03) {
        tracing::error!("不支持的命令或版本: cmd={}, version={}", cmd, version);
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "不支持的命令或版本",
        ));
    }
    if cmd == 0x01 {
        handle_tcp_proxy(stream, buffer, addr_type).await?;
    } else if cmd == 0x03 {
        _handle_udp_proxy(stream).await?;
    }
    Ok(())
}

pub async fn handle_tcp_proxy(
    mut stream: TcpStream,
    buffer: [u8; 256],
    addr_type: u8,
) -> io::Result<()> {
    let (target_addr, target_port) = match addr_type {
        0x01 => {
            // IPv4 地址
            let addr = buffer[4..8].to_vec();
            let port = u16::from_be_bytes([buffer[8], buffer[9]]);
            (
                format!("{}.{}.{}.{}", addr[0], addr[1], addr[2], addr[3]),
                port,
            )
        }
        0x03 => {
            // 域名
            let len = buffer[4] as usize;
            let addr = String::from_utf8_lossy(&buffer[5..5 + len]).to_string();
            let port = u16::from_be_bytes([buffer[5 + len], buffer[6 + len]]);
            (addr, port)
        }
        _ => {
            tracing::error!("不支持的地址类型: {}", addr_type);
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "不支持的地址类型",
            ));
        }
    };
    tracing::debug!("目标地址TCP: {}:{}", target_addr, target_port);

    // 3. 建立目标连接
    // println!("建立目标[{}:{}]连接...", target_addr, target_port);
    let target_stream = TcpStream::connect((target_addr.as_str(), target_port)).await?;
    // println!("建立目标连接成功");
    let (mut target_reader, mut target_writer) = io::split(target_stream);

    // 4. 发送成功响应
    // tracing::debug!("发送握手成功...{}:{}", target_addr, target_port);
    stream
        .write_all(&[0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
        .await?;
    // tracing::debug!("开启数据转发: {}:{}", target_addr, target_port);

    // 5. 数据转发
    // 把stream分割为reader与writer
    let (mut client_reader, mut client_writer) = io::split(stream);

    let client_to_target = async {
        io::copy(&mut client_reader, &mut target_writer).await?;
        target_writer.shutdown().await?;
        Ok::<_, io::Error>(())
    };

    let target_to_client = async {
        io::copy(&mut target_reader, &mut client_writer).await?;
        // client_writer.shutdown().await?;
        Ok::<_, io::Error>(())
    };
    tokio::try_join!(client_to_target, target_to_client)?;
    // println!("完成转发");
    Ok(())
}

pub async fn _handle_udp_proxy(mut stream: TcpStream) -> io::Result<()> {
    // 绑定 UDP 套接字（端口由系统自动分配）
    let udp_socket = UdpSocket::bind("0.0.0.0:0").await?;
    let udp_port = udp_socket.local_addr()?.port();
    // 返回 UDP Associate 响应（格式：VER REP RSV ATYP BND.ADDR BND.PORT）
    let response = [
        0x05,
        0x00,
        0x00, // VER REP RSV
        0x01, // ATYP=IPv4
        0x00,
        0x00,
        0x00,
        0x00, // BND.ADDR=0.0.0.0
        (udp_port >> 8) as u8,
        (udp_port & 0xFF) as u8, // BND.PORT
    ];
    stream.write_all(&response).await?;
    // 启动 UDP 转发循环（见步骤三）
    tokio::spawn(udp_relay(udp_socket));
    // 保持 TCP 连接存活（UDP Associate 需要 TCP 连接保持）
    loop {
        if stream.readable().await.is_err() {
            break;
        }
        tokio::time::sleep(Duration::from_secs(300)).await;
    }
    Ok(())
}
pub async fn udp_relay(socket: UdpSocket) -> io::Result<()> {
    let mut buf = [0u8; 65536];
    loop {
        let (len, src_addr) = socket.recv_from(&mut buf).await?;

        // 解析 SOCKS5 UDP 包头（格式：RSV FRAG ATYP DST.ADDR DST.PORT DATA）
        if len < 10 || buf[0..2] != [0x00, 0x00] {
            // 检查 RSV 和 FRAG
            continue; // 丢弃无效包
        }

        let (dst_addr, data_start) = match buf[3] {
            // ATYP 位于第3字节
            0x01 => {
                // IPv4
                let ip = Ipv4Addr::from(u32::from_be_bytes(buf[4..8].try_into().unwrap()));
                tracing::debug!("目标地址UDP: {}", ip.to_string());
                (
                    SocketAddr::V4(SocketAddrV4::new(ip, u16::from_be_bytes([buf[8], buf[9]]))),
                    10,
                )
            }
            0x03 => {
                // 域名
                let domain_len = buf[4] as usize;
                let domain = String::from_utf8_lossy(&buf[5..5 + domain_len]);
                let port = u16::from_be_bytes([buf[5 + domain_len], buf[6 + domain_len]]);
                tracing::debug!("目标地址UDP: {}:{}", domain, port);
                (
                    SocketAddr::new(domain.parse().unwrap(), port),
                    7 + domain_len,
                )
            }
            _ => continue, // 不支持的地址类型
        };

        // 转发数据到目标地址
        socket.send_to(&buf[data_start..len], dst_addr).await?;

        // 接收目标返回数据并回传客户端
        let mut resp_buf = [0u8; 65536];
        let resp_len = match socket.recv_from(&mut resp_buf).await {
            Ok((len, _)) => len,
            Err(_) => continue,
        };

        // 构造 SOCKS5 UDP 包头
        let mut header = Vec::with_capacity(10);
        header.extend_from_slice(&[0x00, 0x00, 0x00]); // RSV + FRAG
        header.push(0x01); // ATYP=IPv4
        let tmp_ip = src_addr.ip();

        match src_addr.ip() {
            IpAddr::V4(ip) => {
                header.extend_from_slice(&ip.octets());
            }
            IpAddr::V6(ip) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "不支持的IPv6地址",
                ));
            }
        }
        // header.extend_from_slice(&src_addr.ip().to_ipv4().unwrap().octets());
        header.extend_from_slice(&src_addr.port().to_be_bytes());

        // 组合包头和数据
        let packet = [&header[..], &resp_buf[..resp_len]].concat();
        socket.send_to(&packet, src_addr).await?;
    }
}
